import os
import sys
import glob
import argparse
from pathlib import Path
from typing import List, Optional, Tuple
from dataclasses import dataclass
from pprint import pprint
import cv2
import numpy as np
from PIL import Image

# 全局变量：控制详细日志输出
is_verbose = False

def verbose_print(msg: str):
    if is_verbose:
        print(msg)


class ImageDiffInfo:
    """
    图片差异信息类。
    """
    def __init__(self, img1: np.ndarray, img2: np.ndarray, channel_idx:int, diff_img: np.ndarray, diff_threshold, diff_mask=None,border_mask=None):
        """
        初始化图片差异信息类。
        参数：
        - img1 (np.ndarray): 第一张图片。
        - img2 (np.ndarray): 第二张图片。
        - channel_idx (int): 通道索引，默认为-1（所有通道）。
        - diff_img (np.ndarray): 差异图片。
        - diff_threshold (int): 差分阈值，超过阈值才计算差。
        - diff_mask (np.ndarray): 差异掩码。
        - border_mask (np.ndarray): 边框掩码。
        """
        self.img1 = img1
        self.img2 = img2
        self.channel_idx = channel_idx
        self.diff_img = diff_img    # 差异图片,格式为float32
        self.diff_threshold = diff_threshold
    
    
    
def edge_mask(gray: np.ndarray, dilate_kernel = (3,3)) -> np.ndarray:
    """
    边缘检测和二值化。
    参数：
    - gray (np.ndarray): 输入灰度图片。
    - dilate_kernel (tuple): 膨胀核大小，默认为(3,3)。
    返回值：
    - np.ndarray: 边缘检测mask。
    """
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    v = np.median(blur)
    sigma = 0.33
    canny_low = int(max(0, (1.0 - sigma) * v))
    canny_high = int(min(255, (1.0 + sigma) * v))
    edges = cv2.Canny(blur, canny_low, canny_high)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, dilate_kernel)
    edges_dilated = cv2.dilate(edges, kernel, iterations=1)
    mask_bool = edges_dilated > 0  # dtype=bool
    return mask_bool

def image_compare(img1: np.ndarray, img2: np.ndarray, channel_idx, diff_threshold = 3) -> ImageDiffInfo:
    """
    计算两张图片的相似度。
    参数：
    - img1 (np.ndarray): 第一张图片。
    - img2 (np.ndarray): 第二张图片。
    - channel_idx (int): 通道索引，默认为-1（所有通道）。
    - diff_threshold (int): 差分阈值，超过阈值才计算差分，默认为4。
    返回值：
    - ImageDiffInfo: 包含差异信息的类实例。
    """
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    if channel_idx == -1:
        diff_img = img2 - img1
    else:
        diff_img = img2[:, :, channel_idx] - img1[:, :, channel_idx]
    diff_mask = abs(diff_img) > diff_threshold
    diff_img[diff_mask == False] = 0
    return ImageDiffInfo(img1=img1, img2=img2, channel_idx=channel_idx, diff_img=diff_img, diff_threshold=diff_threshold, diff_mask=diff_mask)


def convert_image_to_rgb(img: np.ndarray, offset=0) -> np.ndarray:
    """
    将图片转换为指定格式。
    参数：
    - img (np.ndarray): 输入图片
    - offset (int): 偏移量，默认为0。
    返回值：
    - np.ndarray: 转换后的图片
    """
    dst = img.copy() + offset
    dst = dst.astype(np.uint8)
    dst = cv2.cvtColor(dst, cv2.COLOR_YUV2RGB) 
    return dst

def calc_enhance_sharpen(gt_img: np.ndarray, tnr_img: np.ndarray, border_mask: np.ndarray)-> np.ndarray:
    """
    计算增强锐化图
    sharpen = gt_img - tnr_img if coord in border_mask
    sharpen = 0 otherwise
    参数：
    - gt_img (np.ndarray): groundtruth参考图片
    - tnr_img (np.ndarray): 长曝光时域降噪增强图片
    - border_mask (np.ndarray): 边框mask
    返回值：
    - np.ndarray: 增强锐化图
    """
    if gt_img is None or tnr_img is None or border_mask is None:
        raise ValueError("gt_img, tnr_img, border_mask不能为空")
    if gt_img.shape != tnr_img.shape:
        raise ValueError("gt_img和border_mask形状不一致")
    if tnr_img.shape[:2] != border_mask.shape[:2]:
        raise ValueError("tnr_img和border_mask形状不一致")
    
    # 创建锐化图，初始化为0
    sharpen = np.zeros_like(gt_img, dtype=np.float32)
    # 在border_mask为True的位置计算gt_img - tnr_img
    sharpen[border_mask] = gt_img[border_mask] - tnr_img[border_mask]
    return sharpen

def calc_noise_image(origin_img: np.ndarray, gt_img: np.ndarray, sharpen_img: np.ndarray)-> np.ndarray:
    """
    计算噪声图
    noise = gt_img - origin_img - (gt_img - sharpen_img)
    参数：
    - origin_img (np.ndarray): 原始图片
    - gt_img (np.ndarray): groundtruth参考图片
    - sharpen_img (np.ndarray): 增强锐化图
    返回值：
    - np.ndarray: 噪声图
    """
    if origin_img is None or gt_img is None or sharpen_img is None:
        raise ValueError("origin_img, gt_img, sharpen_img不能为空")
    if origin_img.shape != gt_img.shape:
        raise ValueError("origin_img和gt_img形状不一致")
    if origin_img.shape != sharpen_img.shape:
        raise ValueError("origin_img和sharpen_img形状不一致")
    
    # 计算噪声图：noise = origin_img - (gt_img - sharpen_img)
    noise = origin_img.astype(np.float32) - (gt_img.astype(np.float32) - sharpen_img)
    return noise


def main():
    """
    主函数：处理命令行参数并执行图片差分分析
    """
    global is_verbose
    
    parser = argparse.ArgumentParser(
        description='图片差分分析工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog="""
使用示例:
        """
    )
    
    # 必需参数
    parser.add_argument('--origin', help='原始图片路径')
    parser.add_argument('--gt', help='groundtruth图片路径')
    parser.add_argument('--tnr', help='时域降噪图片路径')
    
    # 可选参数
    parser.add_argument('-t','--diff_threshold', type=int, default=4,help='差分阈值，超过阈值才计算差分，默认为4')
    parser.add_argument('-o','--output_dir', help='差分图和mask图片保存目录')
    parser.add_argument('-v','--verbose', action='store_true',help='启用详细输出模式')
    
    args = parser.parse_args()
    
    # 设置详细输出模式
    is_verbose = args.verbose
    
    verbose_print("开始图片差分分析...")
    verbose_print(f"原始图片: {args.origin}")
    verbose_print(f"groundtruth图片: {args.gt}")
    verbose_print(f"时域降噪图片: {args.tnr}")
    verbose_print(f"差分阈值: {args.diff_threshold}")
    verbose_print(f"输出目录: {args.output_dir}")
    
    # 加载图片
    verbose_print("正在加载图片...")
    origin_img = np.array(Image.open(args.origin))
    gt_img = np.array(Image.open(args.gt))
    tnr_img = np.array(Image.open(args.tnr))
    origin_img = cv2.cvtColor(origin_img, cv2.COLOR_RGB2YUV)
    gt_img = cv2.cvtColor(gt_img, cv2.COLOR_RGB2YUV)
    tnr_img = cv2.cvtColor(tnr_img, cv2.COLOR_RGB2YUV)
    # 检查图片尺寸
    if origin_img.shape != gt_img.shape or origin_img.shape != tnr_img.shape:
        raise ValueError(f"图片格式不匹配: origin_img.shape:{origin_img.shape} vs origin_img.shape:{gt_img.shape} vs tnr_img.shape:{tnr_img.shape}")
    
    # 检测边框
    border_mask = edge_mask(gt_img[:,:,0], dilate_kernel=(3,3))
    # 执行图片差分
    # diff_info = image_compare(origin_img, gt_img, channel_idx = args.channel, diff_threshold=args.diff_threshold)
    sharpen_img = calc_enhance_sharpen(gt_img=gt_img, tnr_img=tnr_img, border_mask=border_mask)
    noise_img = calc_noise_image(origin_img=origin_img, gt_img=gt_img, sharpen_img=sharpen_img)
    # 保存差分图
    if args.output_dir is None:
        raise ValueError("未指定输出目录")
    dir_path = Path(args.output_dir)
    if not dir_path.exists():
        os.makedirs(dir_path)
    origin_name = Path(args.origin).stem
    # 根据指定的输出格式转换图片
    sharpen_path = dir_path / f"{origin_name}_sharpen.png"
    noise_path = dir_path / f"{origin_name}_noise.png"
    sharpen_rgb_img = convert_image_to_rgb(sharpen_img, offset=128)
    Image.fromarray(sharpen_rgb_img).save(str(sharpen_path))
    noise_rgb_img = convert_image_to_rgb(noise_img, offset=128)
    Image.fromarray(noise_rgb_img).save(str(noise_path))
    verbose_print(f"保存noise图到: {noise_path}")
    verbose_print(f"保存sharpen图到: {sharpen_path}")
    # 数据验证
    denoise_path = dir_path / f"{origin_name}_denoise.png"
    enhance_path = dir_path / f"{origin_name}_enhance.png"
    denoise_img = origin_img - noise_img
    enhance_img = denoise_img + sharpen_img
    denoise_rgb_img = convert_image_to_rgb(denoise_img)
    Image.fromarray(denoise_rgb_img).save(str(denoise_path))
    verbose_print(f"保存denoise图到: {denoise_path}")
    enhance_rgb_img = convert_image_to_rgb(enhance_img)
    Image.fromarray(enhance_rgb_img).save(str(enhance_path))
    verbose_print(f"保存enhance图到: {enhance_path}")


if __name__ == "__main__":
    main()